package com.winit.core.server.linux;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import com.winit.core.server.base.CmdOutputHandler;
import com.winit.core.server.base.SSHServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import java.io.File;


/**
 * 抽象一台可以执行shell指令的linux服务器.
 * <p>
 * Created by keel2008 on 16/8/9.
 */
public class LinuxSSHServer implements SSHServer {

    private static final Logger logger = LoggerFactory.getLogger(LinuxSSHServer.class);

    protected final JSch jsch = new JSch();

    protected Session session;

    protected LoginInfo loginInfo;

    //Timeout 毫秒级,默认5秒.
    protected Integer connectionTimeout = 5000;

    //默认的消息处理器,只是输出到控制台.
    protected CmdOutputHandler cmdOutputHandler = new CmdOutputHandler() {
        @Override
        public void onMessage(String cmd, int index, String message,Boolean isError) {
            logger.info("第" + index + "行命令\t[ " + cmd + " ]\t输出: " + message);
        }
    };

    public LinuxSSHServer(String ip, String userName, Integer port) throws Exception {

        LoginInfo loginInfo = new LoginInfo();
        loginInfo.ip = ip;
        loginInfo.userName = userName;
        loginInfo.port = port;

        this.init(loginInfo);

        try {
            jsch.addIdentity("~/.ssh/id_rsa");
        }catch (Exception e){
            // 无法add就算了,有可能这个默认的私钥是不存在的.
            logger.warn("默认的私钥不存在,将被忽略",e.getMessage());
        }

    }

    public void setPassword(String password) {
        this.loginInfo.password = password;
    }

    public void addPrivateKeyPath(String path) throws Exception {

        File pk = new File(path);

        if (!pk.exists() || pk.isDirectory()) {
            throw new IllegalArgumentException("Pubkey file not found:" + path);
        }

        jsch.addIdentity(path);
    }

    protected void init(LoginInfo loginInfo) throws Exception {

        this.loginInfo = loginInfo;

        this.session = jsch.getSession(loginInfo.userName, loginInfo.ip, loginInfo.port);
    }


    public void setConnectionTimeout(Integer connectionTimeout) {
        this.connectionTimeout = connectionTimeout;
    }

    /**
     * 创建连接,
     *
     * @throws Exception
     */
    @Override
    public void connect() throws Exception {

        if (session != null && session.isConnected()) {
            return;
        } else if (session == null) {
            this.init(loginInfo);
        }

        doConnect(session, loginInfo, connectionTimeout);
    }

    protected static void doConnect(Session _session, LoginInfo _loginInfo, Integer _connectionTimeout) throws Exception {
        java.util.Properties config = new java.util.Properties();
        config.put("StrictHostKeyChecking", "no");
        _session.setConfig(config);

        if (_loginInfo.password != null) {
            _session.setPassword(_loginInfo.password);
        }

        _session.connect(_connectionTimeout);

        logger.info("[[Server]]通过SSH登录服务器"+_loginInfo.ip+" port:"+_loginInfo.port);
    }

    /**
     * 断开连接
     *
     * @throws Exception
     */
    @Override
    public void disConnect() throws Exception {
        if (session == null) {
            return;
        }
        //一旦断开连接,无法重连,只能废弃
        session.disconnect();

        logger.info("[[Server]]断开服务器的连接"+loginInfo.ip+" port:"+loginInfo.port);

        this.session = null;
    }


    public LoginInfo getLoginInfo() {
        return loginInfo;
    }

    //配置外部的消息处理器.
    public void setCmdOutputHandler(CmdOutputHandler cmdOutputHandler) {
        this.cmdOutputHandler = cmdOutputHandler;
    }

    /**
     * 执行多条命令.
     *
     * @param cmds
     * @throws Exception
     */
    @Override
    public void doExec(String... cmds) throws Exception {

        this.doExec(null,cmds);

    }

    /**
     * 构造linux命令行.
     * @param subCmds
     * @return
     */
    public String shell(String ... subCmds){
        StringBuilder cmdBuilder = new StringBuilder();

        for (String subCmd:
             subCmds) {
            cmdBuilder.append(subCmd);
        }

        return cmdBuilder.toString();
    }

    /**
     * 按照指令,顺序执行方法,每完成一条指令的输出,才会开始执行下一条指令.
     *
     * @param handler
     * @param cmds
     * @throws Exception
     */
    @Override
    public void doExec(CmdOutputHandler handler, String... cmds) throws Exception {

        Assert.notEmpty(cmds);

        CmdOutputHandler currentHandler = handler == null ? this.cmdOutputHandler : handler;

        for (int i = 0; i < cmds.length; i++) {
            String cmd = cmds[i];

            ChannelExec channel = (ChannelExec) this.session.openChannel("exec");

            try {
                (channel).setCommand(cmd);
                channel.connect();
                channel.setInputStream(null);

                // 将消息的输出,发给消息处理类.
                currentHandler.handleCmdOutput(cmd, i, channel.getInputStream(),channel.getErrStream());

            } catch (Exception e) {
                logger.error("执行第"+i+"条命令:"+cmd+"时抛出异常",e);
                throw e;
            } finally {
                if (channel != null && channel.isConnected()){
                    channel.disconnect();
                }
            }
        }

    }


    protected class LoginInfo {

        protected String userName;
        protected String password;
        protected String ip;
        protected int port;

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append(userName).append(" ").append(password).append(" ").append(ip).append(" ").append(port);
            return sb.toString();
        }
    }

    @Override
    public String toString() {
        return super.toString();
    }
}
